"""
LangGraph 流式处理学习示例

流式处理（Streaming）允许我们逐步获取模型生成的内容，而不是等待整个响应完成。
这在以下场景特别有用：
- 实时显示生成内容（如聊天应用）
- 处理长文本时减少等待时间
- 监控模型的思考过程

LangGraph 支持多种流式模式，主要的有：
1. values - 流式输出最终的状态值
2. updates - 流式输出图的更新过程（节点执行状态）
3. messages - 只流式输出消息
"""

# 示例：streaming_complete_state.py
# 导入异步IO模块 - 用于处理异步操作
import asyncio
import os
# 从类型提示模块中导入Literal - 用于定义有限的字符串选项
from typing import Literal

from dotenv import load_dotenv
# 导入工具装饰器 - 用于将函数转换为AI可调用的工具
from langchain_core.tools import tool
# 导入OpenAI聊天模型
from langchain_openai import ChatOpenAI
# 导入创建react代理的函数 - React模式让模型能够思考并使用工具
from langgraph.prebuilt import create_react_agent

# 加载环境变量（从.env文件）
load_dotenv()


# 使用@tool装饰器定义一个工具函数
@tool
def get_weather(city: Literal["nyc", "sf"]):
    """
    获取指定城市的天气信息

    Args:
        city: 城市名称，只能是"nyc"（纽约）或"sf"（旧金山）

    Returns:
        str: 天气描述字符串
    """
    # 如果城市是纽约
    if city == "nyc":
        return "It might be cloudy in nyc"
    # 如果城市是旧金山
    elif city == "sf":
        return "It's always sunny in sf"
    # 如果城市不是nyc或sf
    else:
        raise AssertionError("Unknown city")


# 将get_weather函数放入工具列表，供代理使用
tools = [get_weather]

# 使用ModelScope API
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="deepseek-ai/DeepSeek-R1-0528",
    streaming=True,  # 关键参数：启用流式处理
    temperature=0.7  # 创造性程度，0-1，越高越有创造性
)

# 使用模型和工具创建一个react代理
# React代理能够让模型思考并决定何时使用工具
graph = create_react_agent(llm, tools)


async def main():
    """
    主异步函数，演示LangGraph的流式处理功能
    """
    # 创建反应式代理
    graph = create_react_agent(llm, tools)

    # 定义输入消息
    inputs = {"messages": [("human", "what's the weather in sf")]}

    print("=== 开始流式处理（updates模式）===")
    print("此模式显示图中每个节点的更新过程\n")

    # 异步迭代获取代理的更新流
    # stream_mode="updates" 会显示每个节点的执行状态变化
    async for chunk in graph.astream(inputs, stream_mode="updates"):
        # 遍历每个更新块中的节点和值
        for node, values in chunk.items():
            print(f"📡 接收到节点更新: '{node}'")
            print(f"   值: {values}")
            print("-" * 50)


async def demo_values_mode():
    """
    演示values流模式的示例函数
    values模式只输出最终的状态值
    """
    print("\n=== 流式处理（values模式）===")
    print("此模式只显示最终的状态值变化\n")

    graph = create_react_agent(llm, tools)
    inputs = {"messages": [("human", "what's the weather in nyc")]}

    async for chunk in graph.astream(inputs, stream_mode="values"):
        for node, values in chunk.items():
            print(f"🎯 最终状态更新 - 节点: '{node}'")
            print(f"   状态值: {values}")
            print("-" * 50)


async def demo_messages_mode():
    """
    演示messages流模式的示例函数
    messages模式只输出消息内容
    """
    print("\n=== 流式处理（messages模式）===")
    print("此模式只显示消息内容的变化\n")

    graph = create_react_agent(llm, tools)
    inputs = {"messages": [("human", "tell me about both nyc and sf weather")]}

    async for chunk in graph.astream(inputs, stream_mode="messages"):
        print(f"💬 新消息: {chunk}")
        print("-" * 50)


if __name__ == "__main__":
    # 运行主要的流式处理示例
    # asyncio.run(main())

    # 取消注释下面的行来体验不同的流模式
    # asyncio.run(demo_values_mode())
    asyncio.run(demo_messages_mode())
